# -*- coding: utf-8 -*-

# 导入pyspark
from pyspark import SparkContext
from pyspark.mllib.classification import NaiveBayes
from pyspark.mllib.regression import LabeledPoint
from pyspark.mllib.linalg import Vectors
from pyspark.mllib.tree import DecisionTree
import numpy as np

# 导入可视化
import matplotlib
matplotlib.use('Agg') # 不回显
import matplotlib.pyplot as plt

# 计时工具
import time

def splitData(data_with_idx, percent):
	testRecords = data_with_idx.sample(False, percent, 42)
	trainRecords = data_with_idx.subtractByKey(testRecords)
	testData = testRecords.map(lambda (idx, p): p)
	trainData = trainRecords.map(lambda (idx, p): p)
	return trainData, testData

# 朴素贝叶斯分类模型
def naiveBayes(trainData, testData):
	startTime = time.time()
	model = NaiveBayes.train(trainData, 1.0)
	predictionAndLabel = testData.map(lambda p: (model.predict(p.features), p.label))
	accuracy = 1.0 * predictionAndLabel.filter(lambda (x, v): x == v).count() / testData.count()
	endTime = time.time()
	durationTime = endTime - startTime
	return durationTime

# 画图
def drawFigure(xArr, yArr, title):
	figName = "%s.png"%title
	plt.figure(num=1, figsize=(8,6))
	plt.title(title, size=14)
	plt.xlabel('x', size=14)
	plt.ylabel('y', size=14)
	plt.plot(xArr, yArr)
	plt.savefig(figName, format='png')

# main函数部分
sc = SparkContext("yarn-client", "NaiveBayes Spark App")

# 加载数据集 
dataPath = "hdfs://192.168.119.141:9100/data/classification/mat_classify"
# 原始数据集
raw_data = sc.textFile(dataPath)
# LabeledPoint数据集
records = raw_data.map(lambda x: x.split(","))
labeledRecords = records.map(lambda x: LabeledPoint(float(x[0]), Vectors.dense([float(y) for y in x[1].split(" ")])))
data_with_idx = labeledRecords.zipWithIndex().map(lambda (k, v): (v, k))

param = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
metrics = []

for percent in param:
	# 拆分训练集与测试集
	trainData, testData = splitData(data_with_idx, percent)
	# 训练朴素贝叶斯分类模型
	durationTime = naiveBayes(trainData, testData)
	metrics.append(durationTime)

drawFigure(param, metrics, 'TimeMeasureNaiveBayes')

# 关闭sc
sc.stop()